# -*- coding: utf-8 -*- #
# -----------------------------------------------------------------------
# File Name:    inference.py
# Version:      ver1_0
# Created:      2024/06/17
# Description:  本文件定义了模型推理功能，并实现GUI界面用于图片上传和结果显示
# -----------------------------------------------------------------------
import torch
from PIL import Image
import sys
from PyQt5.QtWidgets import (QApplication, QMainWindow, QLabel, QPushButton,
                             QFileDialog, QVBoxLayout, QHBoxLayout, QWidget, QMessageBox)
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import Qt
from torchvision.transforms import ToTensor


class InferenceGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.model = None
        self.device = None
        self.init_ui()
        self.load_model()

    def init_ui(self):
        """初始化GUI界面"""
        self.setWindowTitle('手势数字识别系统')
        self.setGeometry(100, 100, 800, 600)

        # 主布局
        main_layout = QHBoxLayout()

        # 左侧图片显示区
        left_widget = QWidget()
        left_layout = QVBoxLayout()
        self.image_label = QLabel("未选择图片")
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setMinimumSize(400, 400)
        self.image_label.setStyleSheet("border: 1px solid #ccc;")
        left_layout.addWidget(self.image_label)

        upload_btn = QPushButton("上传图片")
        upload_btn.clicked.connect(self.upload_image)
        left_layout.addWidget(upload_btn, alignment=Qt.AlignCenter)
        left_widget.setLayout(left_layout)

        # 右侧结果显示区
        right_widget = QWidget()
        right_layout = QVBoxLayout()

        result_label = QLabel("识别结果:")
        result_label.setStyleSheet("font-size: 16px; font-weight: bold;")
        right_layout.addWidget(result_label)

        self.result_display = QLabel("等待识别...")
        self.result_display.setAlignment(Qt.AlignCenter)
        self.result_display.setMinimumSize(300, 100)
        self.result_display.setStyleSheet("font-size: 48px; color: #333;")
        right_layout.addWidget(self.result_display)

        right_widget.setLayout(right_layout)

        # 添加到主布局
        main_layout.addWidget(left_widget)
        main_layout.addWidget(right_widget)

        # 设置中央部件
        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)

    def load_model(self):
        """加载训练好的模型"""
        try:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            self.model = torch.load('./models/model.pkl', map_location=self.device, weights_only=False)
            self.model.to(self.device)
            self.model.eval()
        except Exception as e:
            QMessageBox.critical(self, "错误", f"模型加载失败: {str(e)}")
            self.close()

    def upload_image(self):
        """上传图片并执行推理"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "图片文件 (*.png *.jpg *.jpeg *.bmp)"
        )

        if file_path:
            try:
                # 显示图片
                pixmap = QPixmap(file_path)
                self.image_label.setPixmap(pixmap.scaled(
                    self.image_label.width(), self.image_label.height(),
                    Qt.KeepAspectRatio, Qt.SmoothTransformation
                ))

                # 执行推理
                result = self.inference(file_path)
                self.result_display.setText(str(result))

            except Exception as e:
                QMessageBox.critical(self, "错误", f"图片处理失败: {str(e)}")

    def inference(self, image_path):
        """模型推理核心函数"""
        # 加载并预处理图片
        image = Image.open(image_path).convert('RGB')
        transform = ToTensor()
        image_tensor = transform(image).unsqueeze(0).to(self.device)

        # 前向传播
        with torch.no_grad():
            output = self.model(image_tensor)
            _, predicted = torch.max(output, 1)
            return predicted.item()


def main():
    """主函数"""
    app = QApplication(sys.argv)
    window = InferenceGUI()
    window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()